from typing import List
from pymongo.collection import Collection
from pymongo import IndexModel, DESCENDING, ASCENDING
from datetime import datetime

from ..interface import Job, JobStore, NewJobModel, JobStateType, DONE


class MongodbStore(JobStore):
    '''mongodb存储实现'''
    collection: Collection

    def __init__(self, collection):
        self.collection = collection
        idx1 = IndexModel([('job_id', DESCENDING)], unique=True)
        idx2 = IndexModel([
            ('state', ASCENDING),
            ('update_time', DESCENDING)
        ])
        idx3 = IndexModel([('func_name', DESCENDING)])
        self.collection.create_indexes([idx1, idx2, idx3])

    def get_job(self, job_id: str) -> Job:
        if not isinstance(job_id, str):
            job_id = str(job_id)
        doc = self.collection.find_one({'job_id': job_id})
        return Job(**doc)

    def create_job(self, job: NewJobModel) -> Job:
        self.collection.insert_one(job.get_data())
        return self.get_job(job.job_id)

    def _current_timestamp(self) -> datetime:
        return datetime.now()

    def update_job_state(self, job: Job, state: JobStateType) -> Job:
        if not isinstance(state, JobStateType):
            raise TypeError('state type must be an JobStateType')
        update_time = self._current_timestamp()
        res = self.collection.update_one(
            {
                'job_id': job.job_id,
                'sched_times': job.sched_times
            },
            {
                '$set': {
                    'state': state,
                    'sched_times': job.sched_times + 1,
                    'update_time': update_time.timestamp()
                }
            }
        )
        if res.modified_count == 1:
            job.sched_times = job.sched_times + 1
            job.update_time = update_time
            job.state = state
        return Job

    def load_all_jobs(self) -> List[Job]:
        docs = self.collection.find(
            {'state': {'$ne': DONE}}
        )
        jobs = []
        for doc in docs:
            jobs.append(Job(**doc))
        return jobs


if __name__ == '__main__':
    import random
    import time
    from pymongo import MongoClient
    from ..interface import RUNNING
    cli = MongoClient()
    db = cli['jobs']
    db.drop_collection('demo')
    collection = db.get_collection('demo')
    store = MongodbStore(collection)
    job_id = f'demo_{int(time.time()) + random.randrange(1, 100)}'
    new_job = NewJobModel(
        job_id,
        'demo',
        {'a': 1}
    )
    job = store.create_job(new_job)
    print(job)
    store.update_job_state(job, RUNNING)
    print(job)
    job_id = f'demo_{int(time.time()) + random.randrange(1, 100)}'
    new_job = NewJobModel(
        job_id,
        'demo',
        {'b': 1}
    )
    store.create_job(new_job)
    jobs = store.load_all_jobs()
    print(jobs)
